#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
PDF文档解析器模块
专门解析PDF文档，匹配需求文档。
"""
import pdfplumber
import re
from collections import Counter
import time
import logging
from typing import Dict, List, Any
import os

logging.basicConfig(level=logging.INFO)

class PdfParser:
    def parse_document(self, file_path: str) -> Dict[str, Any]:
        start_time = time.time()
        with pdfplumber.open(file_path) as pdf:
            page_count = len(pdf.pages)
            text_lines = []
            font_info = []
            tables = []
            
            for page_num, page in enumerate(pdf.pages):
                # 提取文本内容（作为后备方案）
                page_text = page.extract_text()
                if page_text:
                    for line_text in page_text.split('\n'):
                        if line_text.strip():
                            # 创建一个简单的文本行对象
                            text_lines.append({
                                'text': line_text.strip(),
                                'size': 12,  # 默认字号
                                'fontname': 'Unknown',
                                'top': 0,
                                'bottom': 0
                            })
                
                # 尝试提取带格式的文本行
                try:
                    extracted_lines = page.extract_text_lines()
                    for line in extracted_lines:
                        if 'text' in line and line['text'].strip():
                            # 计算该行的平均字号
                            if 'chars' in line and line['chars']:
                                sizes = [char.get('size', 12) for char in line['chars'] if 'size' in char]
                                if sizes:
                                    avg_size = sum(sizes) / len(sizes)
                                    line['size'] = avg_size
                                    
                                    # 收集字体信息
                                    for char in line['chars']:
                                        if 'size' in char and 'fontname' in char:
                                            font_info.append((char['size'], 'bold' in char['fontname'].lower()))
                                else:
                                    line['size'] = 12
                            else:
                                line['size'] = 12
                            
                            text_lines.append(line)
                except Exception as e:
                    logging.warning(f"提取带格式文本失败，使用纯文本: {e}")
                    continue
                
                # 提取表格
                try:
                    page_tables = page.extract_tables()
                    for tbl in page_tables:
                        tables.append({
                            "type": "table",
                            "rows": tbl,
                            "source_confidence": "low",
                            "note": "Extracted from PDF, may contain errors.",
                            "page": page_num + 1
                        })
                except Exception as e:
                    logging.warning(f"提取表格失败: {e}")
                    continue
            
            # 如果没有收集到字体信息，使用默认值
            if not font_info:
                base_size = 12
            else:
                # 计算基准字号
                sizes = [info[0] for info in font_info if info[0] > 0]
                base_size = Counter(sizes).most_common(1)[0][0] if sizes else 12
            
            # 如果没有提取到带格式的文本行，创建基本结构
            if not text_lines:
                sections = [{
                    "title": "文档内容",
                    "level": 1,
                    "paragraph": [],
                    "sections": []
                }]
            else:
                # 识别标题等
                sections = self._build_structure(text_lines, base_size)
            
            # 将表格插入到sections中
            for table in tables:
                if sections:
                    # 将表格转换为文本表示
                    rows = table.get("rows", [])
                    table_text = "\n".join(["\t".join(row) for row in rows])
                    sections[-1]["paragraph"].append(f"[表格]\n{table_text}")
                else:
                    rows = table.get("rows", [])
                    table_text = "\n".join(["\t".join(row) for row in rows])
                    sections.append({
                        "title": "内容",
                        "level": 1,
                        "paragraph": [f"[表格]\n{table_text}"],
                        "sections": []
                    })
        
        processing_time = time.time() - start_time
        title = os.path.splitext(os.path.basename(file_path))[0]
        
        return {
            "title": title,
            "sections": sections,
            "metadata": {
                "source_filename": os.path.basename(file_path),
                "page_count": page_count,
                "processing_time_seconds": processing_time
            }
        }

    def _is_title(self, text, size, base_size, is_bold) -> tuple[bool, int]:
        """
        判断文本行是否为标题，并返回标题层级
        """
        # 首先检查是否为图注，如果是则不是标题
        figure_caption_patterns = [
            r'^图\s*\d+',  # 图 1、图1、图 2-1 等
            r'^图\s*[一二三四五六七八九十\d]+',  # 图一、图二等
            r'^Figure\s*\d+',  # Figure 1, Figure 2等
            r'^Fig\s*\.\s*\d+',  # Fig. 1, Fig. 2等
        ]
        
        for pattern in figure_caption_patterns:
            if re.match(pattern, text):
                return False, 0

        # 检查是否为注释，如果是则不是标题
        if re.match(r'^\([^)]*注[^)]*\)', text) or re.match(r'^注\s*[:：]', text):
            return False, 0

        # 新增规则：检查中文标号+内容形式（如"一、找到校正程序"），直接识别为标题
        if re.match(r'^[一二三四五六七八九十]+[、．\.]\s*.*', text):
            return True, 1

        # 检查是否有标号（优先于字号检查）
        level = 0
        
        # 一级标题：第一章、一、1 引言、1. 引言、1.引言
        if re.match(r'^第[一二三四五六七八九十\d]+章', text) or re.match(r'^[一二三四五六七八九十]+[、．\.]', text) or \
           re.match(r'^\d+\s+[^\d\s。，；：！？]{1,15}$', text) or re.match(r'^\d+\.\s+[^\d\s。，；：！？]{1,15}$', text) or re.match(r'^\d+\.\s*[^\d\s。，；：！？]{1,15}$', text):  # 匹配"1 引言"、"1. 引言"和"1.引言"这样的格式，限制标题长度和标点符号
            level = 1
        # 多级数字标题：1.1、1.1.1、1.1.1.1等（直接识别为标题）
        elif re.match(r'^\d+\.\d+', text) and re.match(r'^\d+\.\d+.*[^\d\s。，；：！？]{1,15}$', text):  # 匹配"1.1"及更复杂的格式，限制标题长度和标点符号
            if re.match(r'^\d+\.\d+\.\d+\.\d+\.\d+\.\d+', text):  # 1.1.1.1.1.1
                level = 6
            elif re.match(r'^\d+\.\d+\.\d+\.\d+\.\d+', text):  # 1.1.1.1.1
                level = 5
            elif re.match(r'^\d+\.\d+\.\d+\.\d+', text):  # 1.1.1.1
                level = 4
            elif re.match(r'^\d+\.\d+\.\d+', text):  # 1.1.1
                level = 3
            else:  # 1.1
                level = 2
        # 单个数字+点号：1.、2.等（需要通过字号判断）
        elif re.match(r'^\d+[、．\.]\s*$', text):  # 匹配"1."、"2."等格式
            # 需要通过字号判断是否为标题
            size_ratio = size / base_size if base_size > 0 else 1
            if size_ratio > 1.0 or is_bold:
                level = 1
            else:
                return False, 0
        
        # 如果有标号，直接返回
        if level > 0:
            return True, level
        
        # 无标号：需要通过字号判断
        size_ratio = size / base_size if base_size > 0 else 1
        
        # 更严格的视觉规则判断：字号要大且文本要短且不包含标点符号
        if size_ratio > 1.3 and len(text) <= 20 and not re.search(r'[。，；：！？]', text):
            return True, 0  # 返回0表示需要在_build_structure中动态确定层级
        else:
            return False, 0

    def _is_figure_or_table_caption(self, text) -> bool:
        """
        检查是否为图注或表格标题
        """
        figure_caption_patterns = [
            r'^图\s*\d+',  # 图 1、图1、图 2-1 等
            r'^图\s*[一二三四五六七八九十\d]+',  # 图一、图二等
            r'^Figure\s*\d+',  # Figure 1, Figure 2等
            r'^Fig\s*\.\s*\d+',  # Fig. 1, Fig. 2等
        ]

        table_caption_patterns = [
            r'^表\s*\d+',  # 表 1、表1、表 2-1 等
            r'^表\s*[一二三四五六七八九十\d]+',  # 表一、表二等
            r'^Table\s*\d+',  # Table 1, Table 2等
            r'^Tab\s*\.\s*\d+',  # Tab. 1, Tab. 2等
        ]

        is_figure = any(re.match(pattern, text) for pattern in figure_caption_patterns)
        is_table_caption = any(re.match(pattern, text) for pattern in table_caption_patterns)
        
        return is_figure or is_table_caption

    def _is_list_item(self, text) -> bool:
        """
        检查是否为列表项
        """
        return bool(re.match(r'^[-*]\s|^\d+\.\s', text))

    def _should_end_paragraph(self, text, current_paragraphs) -> bool:
        """
        检查是否应该结束当前段落
        """
        return (
            text.endswith(('。', '.', '？', '！', '?', '!', ':', '：')) or 
            len(text) > 100 or  # 长文本自动分段
            (len(current_paragraphs) > 0 and len(" ".join(current_paragraphs + [text])) > 200)
        )

    def _determine_title_level(self, level, last_numbered_level) -> int:
        """
        确定标题层级
        """
        if level > 0:  # 有标号标题
            return level
        
        # 无标号标题，根据上一个有标号标题的层级确定
        if last_numbered_level > 0:
            level = last_numbered_level + 1
        else:
            level = 2  # 默认2级
        
        # 限制最大层级为6
        return min(level, 6)

    def _add_current_content_to_section(self, stack, current_paragraphs, current_list_items):
        """
        将当前段落和列表项添加到当前section
        """
        if current_paragraphs:
            paragraph_text = " ".join(current_paragraphs)
            if paragraph_text.strip() and stack:
                stack[-1]["paragraph"].append(paragraph_text)
        
        if current_list_items and stack:
            # 将列表项作为段落处理
            stack[-1]["paragraph"].extend(current_list_items)

    def _create_new_section(self, text, level):
        """
        创建新的section
        """
        return {
            "title": text,
            "level": level,
            "paragraph": [],
            "sections": []
        }

    def _update_section_stack(self, stack, sections, new_section):
        """
        更新section栈
        """
        # 调整栈
        while stack and stack[-1]["level"] >= new_section["level"]:
            stack.pop()
        
        if stack:
            stack[-1]["sections"].append(new_section)
        else:
            sections.append(new_section)
        
        stack.append(new_section)

    def _build_structure(self, lines, base_size):
        sections = []
        stack = []
        current_paragraphs = []
        current_list_items = []
        last_numbered_level = 0  # 记录上一个有标号标题的层级

        for line in lines:
            text = line['text'].strip()
            if not text: continue

            # 获取字体大小和粗体信息
            size = line.get('size', base_size)
            fontname = line.get('fontname', '')
            is_bold = 'bold' in fontname.lower() or 'heavy' in fontname.lower() or 'black' in fontname.lower()
            
            # 首先检查是否为图注或表格标题
            if self._is_figure_or_table_caption(text):
                if stack:
                    stack[-1]["paragraph"].append(text)
                continue

            # 检查是否为注释
            if re.match(r'^\([^)]*注[^)]*\)', text) or re.match(r'^注\s*[:：]', text):
                if stack:
                    stack[-1]["paragraph"].append(text)
                continue

            # 检查是否为标题
            is_title, level = self._is_title(text, size, base_size, is_bold)

            if is_title:
                # 添加当前段落和列表项到当前section
                self._add_current_content_to_section(stack, current_paragraphs, current_list_items)
                current_paragraphs = []
                current_list_items = []

                # 更新上一个有标号标题的层级
                if level > 0:  # 有标号标题
                    last_numbered_level = level
                
                # 确定标题层级
                level = self._determine_title_level(level, last_numbered_level)
                
                # 确保标题文本不为空
                if not text.strip():
                    continue

                # 创建新section并更新栈
                new_section = self._create_new_section(text, level)
                self._update_section_stack(stack, sections, new_section)
            else:
                # 处理非标题文本
                if self._is_list_item(text):
                    current_list_items.append(text)
                else:
                    if current_list_items:
                        if stack:
                            # 将列表项作为段落处理
                            stack[-1]["paragraph"].extend(current_list_items)
                        current_list_items = []
                    
                    # 检查是否应该结束段落
                    if self._should_end_paragraph(text, current_paragraphs):
                        current_paragraphs.append(text)
                        paragraph_text = " ".join(current_paragraphs)
                        if paragraph_text.strip():
                            if stack:
                                stack[-1]["paragraph"].append(paragraph_text)
                        current_paragraphs = []
                    else:
                        current_paragraphs.append(text)

        # 清理剩余
        if current_paragraphs:
            if stack:
                stack[-1]["paragraph"].append(" ".join(current_paragraphs))
        if current_list_items:
            if stack:
                # 将列表项作为段落处理
                stack[-1]["paragraph"].extend(current_list_items)

        return sections

    def ocr_parse(self, file_path):
        pass  # TODO: OCR integration
